# -*- coding: utf-8 -*-
# @Time    : 2021/7/7 14:46
# @Author  : LuoTianHang=
# ####################train.py 说明##########################
# this script is the core py to train the model
import random

import numpy as np
import torch
from colorama import Fore
from torch import nn
from torch import optim
from torch.backends import cudnn
from torch.utils.data import DataLoader
from tqdm import tqdm

from config import GetConfig
from datalist import MPII
from metric import PCKh
from model import preresnet56
from utils import weights_init

best_acc = 0


class PoseEstimation(object):
    def __init__(self):
        self.args = GetConfig()
        print(f"-----------{self.args.project_name}-------------")

        # region 项目运行配置
        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        self.device = torch.device('cuda' if use_cuda else 'cpu')
        kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}
        # endregion

        # region 数据集dataloader

        with open("data/mpii/data.json") as f:
            data = f.readlines()
        random.shuffle(data)
        train_num = int(len(data) * 0.75)
        self.train_data = data[:train_num]
        self.test_data = data[train_num:]

        self.train_dataset = MPII(self.train_data)
        self.train_dataloader = DataLoader(self.train_dataset,
                                           batch_size=self.args.train_batch_size, shuffle=True, **kwargs)

        self.test_dataset = MPII(self.test_data, mode="test")
        self.test_dataloader = DataLoader(self.test_dataset,
                                          batch_size=self.args.test_batch_size, shuffle=True, **kwargs)

        self.model = preresnet56(num_classes=16)
        if use_cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count())).to(self.device)
            cudnn.enabled = True
            cudnn.benchmark = True

        # ToDo : use resume or not
        if self.args.resume:
            print("load the weight from pretrained-weight file")
            model_dict = self.model.state_dict()
            pretrained_dict = torch.load(self.args.pretrained_weight, map_location=self.device)['model_state_dict']
            pretrained_dict = {k: v for k, v in pretrained_dict.items() if np.shape(model_dict[k]) == np.shape(v)}
            model_dict.update(pretrained_dict)
            self.model.load_state_dict(model_dict)
            print("Finished to load the weight")
        else:
            print("train from scratch")
            weights_init(self.model, init_type="normal", init_gain=0.1)

        # region 优化器 训练策略 loss定义
        self.optimizer = optim.Adam(self.model.parameters(), self.args.lr, weight_decay=0)
        self.scheduler = optim.lr_scheduler.StepLR(self.optimizer, 3, gamma=0.65)
        # self.scheduler = optim.lr_scheduler.CosineAnnealingWarmRestarts(self.optimizer, T_0=100, T_mult=1)

        self.criterion = nn.MSELoss()
        # endregion

        self.metric = PCKh()

    def start_work(self):
        for epoch in range(1, self.args.epochs):
            self.train(epoch)
            self.test(epoch)

        torch.cuda.empty_cache()
        print("finish model training")

    def train(self, epoch):
        self.model.train()
        pbar = tqdm(self.train_dataloader, desc=f'Test Epoch {epoch}/{self.args.epochs}',
                    bar_format='{l_bar}%s{bar}%s{r_bar}' % (Fore.BLUE, Fore.RESET))
        total_loss = []
        # self.metric.reset()
        for data, target, vis, scale in pbar:
            data, target, vis = data.to(self.device), target.to(self.device), vis.to(self.device)
            self.optimizer.zero_grad()
            outputs = self.model(data)
            # self.metric.update(predictions=outputs.clone().detach(), targets=target.clone(), vis=vis.clone(), scale=scale.clone())
            loss = self.criterion(outputs, target)
            loss.backward()
            self.optimizer.step()
            total_loss.append(loss.item())

            pbar.set_description(
                'Train Epoch :{}/{}'.format(epoch, self.args.epochs) +
                '\t lr:{:22}'.format(self.optimizer.param_groups[0]["lr"]) +
                '\t loss:{:8}'.format(np.around(np.mean(total_loss), 6))
            )
        self.scheduler.step()

    @torch.no_grad()
    def test(self, epoch):
        self.model.eval()
        pbar = tqdm(self.test_dataloader, desc=f'Test Epoch {epoch}/{self.args.epochs}',
                    bar_format='{l_bar}%s{bar}%s{r_bar}' % (Fore.GREEN, Fore.RESET))
        self.metric.reset()
        for data, target, vis, scale in pbar:
            data, target, vis = data.to(self.device), target.to(self.device), vis.to(self.device)
            outputs = self.model(data)
            self.metric.update(predictions=outputs, targets=target, vis=vis, scale=scale)
            pbar.set_description(
                'Test Epoch :{}/{}'.format(epoch, self.args.epochs) +
                '\t Metrics: ' + str(self.metric.get())
            )
        total_acc = self.metric.get()["total"]

        global best_acc
        if total_acc > best_acc:
            best_acc = total_acc
            torch.save({"model_state_dict": self.model.state_dict()}, "weight/best.pth")
            print("model saved")


if __name__ == "__main__":
    model = PoseEstimation()
    model.start_work()
